#ifndef FLANN_WARPER_H_
#define FLANN_WARPER_H_

#include "../Common/GlobalDefines.h"
#include "Eigen/Dense"
#include "flann/flann.hpp"
#include "DistanceFunction.h"
#include <Tools/ZLogManager.h>
#include <Tools/MemoryTools.h>

GLOBAL_NAMESPACE_BEGIN

template <typename Distance>
class FLAnnWarper
{
public:
    FLAnnWarper(int indexParams = 4);
    ~FLAnnWarper();

    void init(const Eigen::MatrixXd& data);
    void destroy();

    void setSearchParams(const flann::SearchParams& params);

    // query one feature(point)
    int knnSearch(const Eigen::VectorXd& queryPoint,
                    std::vector<int>& indices,
                    std::vector<double>& dists,
                    size_t knn);
    int radiusSearch(const Eigen::VectorXd& queryPoint,
                    std::vector<int>& indices,
                    std::vector<double>& dists,
                    double radius);

    // query many features(points)
    int knnSearch(const Eigen::MatrixXd& querys, 
                    std::vector< std::vector<int> >& indices, 
                    std::vector< std::vector<double> >& dists, 
                    size_t knn);
    int radiusSearch(const Eigen::MatrixXd& querys, 
                    std::vector< std::vector<int> >& indices, 
                    std::vector< std::vector<double> >& dists, 
                    double radius);

private:
    flann::SearchParams searchParams_;
    flann::IndexParams  indexParams_;

    flann::Index<Distance> *index_;

    double *dataBlock_;
};

//#include "FLAnnWarper.cpp"

template <typename Distance>
FLAnnWarper<Distance>::FLAnnWarper(int indexParam)
    : index_(NULL)
    , indexParams_(flann::KDTreeIndexParams(indexParam))
    //, indexParams_(flann::KDTreeSingleIndexParams(indexParam))
    , searchParams_(flann::SearchParams(128))
    , dataBlock_(NULL)
{

}

template <typename Distance>
FLAnnWarper<Distance>::~FLAnnWarper()
{
    destroy();
}

template <typename Distance>
void FLAnnWarper<Distance>::destroy()
{
    if (dataBlock_) delete [] dataBlock_;
    dataBlock_ = NULL;
    if (index_) delete index_;
    index_ = NULL;
}

template <typename Distance>
void FLAnnWarper<Distance>::setSearchParams(const flann::SearchParams& params)
{
    searchParams_ = params;
}

template <typename Distance>
void FLAnnWarper<Distance>::init(const Eigen::MatrixXd& data)
{
    int nDataSize = data.rows();
    int nDataDim = data.cols();
    //double *dData = new double[nDataSize*nDataDim];
    //Eigen::MatrixXd dataT = data;
    //dataT.transposeInPlace();
    //memcpy(dData, dataT.data(), nDataSize*nDataDim*sizeof(double));
    dataBlock_ = new double[nDataSize*nDataDim];
    MemoryTools::copyMemery(data, dataBlock_);
    flann::Matrix<double> flMatrix(dataBlock_, nDataSize, nDataDim);

    if (index_) delete index_;
    index_ = new flann::Index<Distance>(flMatrix, indexParams_);
    index_->buildIndex();

    //delete[] dData;
}

template <typename Distance>
int FLAnnWarper<Distance>::knnSearch(const Eigen::VectorXd& query, std::vector<int>& indices, std::vector<double>& dists, size_t knn)
{
    int nDataSize = 1;
    int nDataDim = query.rows();
    double *dData = new double[nDataSize*nDataDim];
    MemoryTools::copyMemery(query, dData);
    //memcpy(dData, query.data(), nDataSize*nDataDim*sizeof(double));

    flann::Matrix<double> flMatrix(dData, nDataSize, nDataDim);
    std::vector< std::vector<int> > nnIdx;
    std::vector< std::vector<double> > nnDists;
    int ret = index_->knnSearch(flMatrix, nnIdx, nnDists, knn, searchParams_);
    indices.resize(knn);
    dists.resize(knn);
    for (size_t i = 0; i < knn; i++)
    {
        indices[i] = nnIdx[0][i];
        dists[i] = nnDists[0][i];
    }
    delete[] dData;
    return ret;
}

template <typename Distance>
int FLAnnWarper<Distance>::radiusSearch(const Eigen::VectorXd& query, std::vector<int>& indices, std::vector<double>& dists, double radius)
{
    int nDataSize = 1;
    int nDataDim = query.rows();
    double *dData = new double[nDataSize*nDataDim];
    MemoryTools::copyMemery(query, dData);
    //memcpy(dData, query.data(), nDataSize*nDataDim*sizeof(double));

    flann::Matrix<double> flMatrix(dData, nDataSize, nDataDim);
    std::vector< std::vector<int> > nnIdx;
    std::vector< std::vector<double> > nnDists;
    int ret = index_->radiusSearch(flMatrix, nnIdx, nnDists, radius, searchParams_);
    indices.resize(nnIdx[0].size());
    dists.resize(nnIdx[0].size());
    for (size_t i = 0; i < nnIdx[0].size(); i++)
    {
        indices[i] = nnIdx[0][i];
        dists[i] = nnDists[0][i];
    }
    delete[] dData;
    return ret;
}

template <typename Distance>
int FLAnnWarper<Distance>::knnSearch(const Eigen::MatrixXd& querys, std::vector< std::vector<int> >& indices, std::vector< std::vector<double> >& dists, size_t knn)
{
    int nDataSize = querys.rows();
    int nDataDim = querys.cols();
    double *dData = new double[nDataSize*nDataDim];
    //Eigen::MatrixXd querysT = querys;
    //querysT.transposeInPlace();
    //memcpy(dData, querysT.data(), nDataSize*nDataDim*sizeof(double));
    MemoryTools::copyMemery(querys, dData);
    //DebugLog << dData[0] << "," << dData[1] << "," << dData[2] << "\n";
    //DebugLog << dData[3] << "," << dData[4] << "," << dData[5] << "\n";

    flann::Matrix<double> flMatrix(dData, nDataSize, nDataDim);
    int ret = index_->knnSearch(flMatrix, indices, dists, knn, searchParams_);
    delete[] dData;
    return ret;
}

template <typename Distance>
int FLAnnWarper<Distance>::radiusSearch(const Eigen::MatrixXd& querys, std::vector< std::vector<int> >& indices, std::vector< std::vector<double> >& dists, double radius)
{
    int nDataSize = querys.rows();
    int nDataDim = querys.cols();
    double *dData = new double[nDataSize*nDataDim];
    //Eigen::MatrixXd querysT = querys;
    //querysT.transposeInPlace();
    //memcpy(dData, querysT.data(), nDataSize*nDataDim*sizeof(double));
    MemoryTools::copyMemery(querys, dData);

    flann::Matrix<double> flMatrix(dData, nDataSize, nDataDim);
    int ret = index_->radiusSearch(flMatrix, indices, dists, radius, searchParams_);
    delete[] dData;
    return ret;
}

GLOBAL_NAMESPACE_END


#endif